import torch
import torch.nn as nn
import torch.nn.functional as F

# 定义简单的CNN模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # 卷积层1
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2)
        # 池化层1
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 卷积层2
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, padding=2)
        # 池化层2
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 全连接层1
        self.fc1 = nn.Linear(32 * 7 * 7, 128)  # 输入尺寸需要根据池化后的特征图计算
        # 全连接层2（输出层）
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # 卷积层1 + ReLU激活 + 池化层1
        x = self.pool1(F.relu(self.conv1(x)))
        # 卷积层2 + ReLU激活 + 池化层2
        x = self.pool2(F.relu(self.conv2(x)))
        # 展平特征图
        x = x.view(-1, 32 * 7 * 7)  # 展平为 [batch_size, 32*7*7]
        # 全连接层1 + ReLU激活
        x = F.relu(self.fc1(x))
        # 全连接层2（输出层）
        x = self.fc2(x)
        return x

# 实例化模型
model = SimpleCNN().cuda()
print(model)